#!/usr/bin/env python

import sys
import math
import rospy
from sensor_msgs.msg import Image
from vision_messages.srv import ObjPoseEst
from geometry_msgs.msg import Point
from visualization_msgs.msg import Marker, MarkerArray
import numpy as np
import cv2
import message_filters
# import threading
 
def thread_job():
    rospy.spin()

lines = [[0, 1], [1, 2], [2, 3], [3, 0], [4, 5], [5, 6],
         [6, 7], [7, 4], [0, 4], [1, 5], [2, 6], [3, 7]]


def R2q(R):
    
    w = R[0,0]+R[1,1]+R[2,2]+1
    x = R[0,0]-R[1,1]-R[2,2]+1
    y = -R[0,0]+R[1,1]-R[2,2]+1
    z = -R[0,0]-R[1,1]+R[2,2]+1
    
    q = np.array([w,x,y,z])

    index = np.argmax(q)

    # assert q[index]>0, "max(q) > 0"
    if q[index]<=0:
        return np.array([0,0,0,1],dtype=np.float32)

    q[index] = math.sqrt(q[index]) / 2

    if index==0:
        q0 = q[index]
        q1 = (R[2,1]-R[1,2]) / (4*q0)
        q2 = (R[0,2]-R[2,0]) / (4*q0)
        q3 = (R[1,0]-R[0,1]) / (4*q0)
    elif index==1:
        q1 = q[index]
        q0 = (R[2,1]-R[1,2]) / (4*q1)
        q2 = (R[0,1]+R[1,0]) / (4*q1)
        q3 = (R[2,0]+R[0,2]) / (4*q1)
    elif index==2:
        q2 = q[index]
        q0 = (R[0,2]-R[2,0]) / (4*q2)
        q1 = (R[0,1]+R[1,0]) / (4*q2)
        q3 = (R[1,2]+R[2,1]) / (4*q2)
    elif index==3:
        q3 = q[index]
        q0 = (R[1,0]-R[0,1]) / (4*q3)
        q1 = (R[2,0]+R[0,2]) / (4*q3)
        q2 = (R[1,2]+R[2,1]) / (4*q3)
    else:
        raise ValueError('index error:' +  str(index))
    # print(np.linalg.norm(np.array([q1,q2,q3,q0],dtype=np.float32)))
    return np.array([q1,q2,q3,q0],dtype=np.float32)

def get_3d_bbox(size, shift=0):
    """
    Args:
        size: [3] or scalar
        shift: [3] or scalar
    Returns:
        bbox_3d: [3, N]
    """
    bbox_3d = np.array([[+size[0] / 2, +size[1] / 2, +size[2] / 2],
                        [+size[0] / 2, +size[1] / 2, -size[2] / 2],
                        [-size[0] / 2, +size[1] / 2, +size[2] / 2],
                        [-size[0] / 2, +size[1] / 2, -size[2] / 2],
                        [+size[0] / 2, -size[1] / 2, +size[2] / 2],
                        [+size[0] / 2, -size[1] / 2, -size[2] / 2],
                        [-size[0] / 2, -size[1] / 2, +size[2] / 2],
                        [-size[0] / 2, -size[1] / 2, -size[2] / 2]]) + shift
    bbox_3d = bbox_3d.transpose()
    return bbox_3d

def transform_coordinates_3d(coordinates, sRT):
    """
    Args:
        coordinates: [3, N]
        sRT: [4, 4]
    Returns:
        new_coordinates: [3, N]
    """
    assert coordinates.shape[0] == 3
    coordinates = np.vstack([coordinates, np.ones((1, coordinates.shape[1]), dtype=np.float32)])
    new_coordinates = np.matmul(sRT, coordinates)
    new_coordinates = new_coordinates[:3, :] / new_coordinates[3, :]
    #new_coordinates = new_coordinates[:3, :]
    return new_coordinates

def draw_detections(pred_sRT, pred_size):
    """ Visualize pose predictions.
    """
    transformed_bbox_3d = np.zeros((pred_sRT.shape[0], 3, 8))
    for i in range(pred_sRT.shape[0]):
        sRT = pred_sRT[i, :, :]
        bbox_3d = get_3d_bbox(pred_size[i, :], 0)
        transformed_bbox_3d[i, :, :] = transform_coordinates_3d(bbox_3d, sRT)
        # transformed_bbox_3d[i, :, 8] = np.mean(transformed_bbox_3d[i, :, :8], axis=1)
    return transformed_bbox_3d

class Detector:
    def __init__(self):
        rospy.init_node('detectbbox3d', anonymous=True)
        self.rate = rospy.Rate(5)

        self.marker_pub = rospy.Publisher(
            '/detect_bbox3d', MarkerArray, queue_size=2)

        self.marker_array = MarkerArray()

    def get_3d_box(self, center, box_size, R):
        ''' Calculate 3D bounding box corners from its parameterization.

        Input:heading_angle
            box_size: tuple of (l,w,h)
            : rad scalar, clockwise from pos z axis
            center: tuple of (x,y,z)
        Output:
            corners_3d: numpy array of shape (8,3) for 3D box cornders
        '''
        l, w, h = box_size
        x_corners = [l/2, l/2, -l/2, -l/2, l/2, l/2, -l/2, -l/2]
        y_corners = [w/2, -w/2, -w/2, w/2, w/2, -w/2, -w/2, w/2]
        z_corners = [h/2, h/2, h/2, h/2, -h/2, -h/2, -h/2, -h/2]
        corners_3d = np.dot(R, np.vstack([x_corners, y_corners, z_corners]))
        corners_3d[0, :] = corners_3d[0, :] + center[0]
        corners_3d[1, :] = corners_3d[1, :] + center[1]
        corners_3d[2, :] = corners_3d[2, :] + center[2]
        corners_3d = np.transpose(corners_3d)
        return corners_3d

    def display(self, boxes):
        # self.marker_array.markers.clear()

        for obid in range(len(boxes)):
            ob = boxes[obid]
            tid = 0
            detect_points_set = []
            for i in range(0, 8):
                detect_points_set.append(Point(ob[i], ob[i+8], ob[i+16]))

            marker = Marker()
            marker.header.frame_id = 'camera_color_optical_frame'
            marker.header.stamp = rospy.Time.now()

            marker.id = obid*2
            marker.action = Marker.ADD
            marker.type = Marker.LINE_LIST

            marker.lifetime = rospy.Duration(0)

            marker.color.r = 1
            marker.color.g = 0
            marker.color.b = 0

            marker.color.a = 1
            marker.scale.x = 0.1
            marker.scale.y = 0.1
            marker.scale.z = 0.1
            marker.points = []

            for line in lines:
                marker.points.append(detect_points_set[line[0]])
                marker.points.append(detect_points_set[line[1]])

        self.marker_array.markers.append(marker)

        self.marker_pub.publish(self.marker_array)
        self.rate.sleep()

class Pose_Marker():
    def __init__(self, pose):
        rospy.init_node('posemarker', anonymous=True)
        self.rate = rospy.Rate(20)
        self.marker_pub = rospy.Publisher("/cube", Marker, queue_size=10)
        rospy.loginfo("Initializing...")
        self.pose = pose

    def pose_marker(self):
        marker = Marker()
        marker.header.frame_id = "camera_color_optical_frame"
        
        marker.header.stamp = rospy.Time.now()
        
        marker.ns = "basic_shapes"
        
        marker.id = 0
        
        marker.type = Marker.CYLINDER
        
        marker.scale.x = 0.2
        marker.scale.y = 0.2
        marker.scale.z = 0.4
        
        marker.action = Marker.ADD
        marker.pose.position.x = self.pose[0, 3]
        marker.pose.position.y = self.pose[1, 3]
        marker.pose.position.z = self.pose[2, 3]
        q = R2q(self.pose[:3, :3])
        marker.pose.orientation.x = q[0]
        marker.pose.orientation.y = q[1]
        marker.pose.orientation.z = q[2]
        marker.pose.orientation.w = q[3]

        marker.color.r = 0.0
        marker.color.g = 0.8
        marker.color.b = 0.0
        marker.color.a = 0.5

        marker.lifetime = rospy.Duration()
        while not rospy.is_shutdown():
            self.marker_pub.publish(marker)
            self.rate.sleep()
            rospy.loginfo("Sending pose marker...")

if __name__ == "__main__":
    # while True:
    #     rospy.init_node('percenption_client_debug')
    #     percenption = Perception_Client()
    #     result = percenption.client()
    #     if result is not None:
    #         result = result.obj_pose_list[0]
    #         obj_pose = result.RT
    #         obj_cate = result.category_name
    #         obj_scales = result.scales
    #         obj_scale = result.scale
    #         rospy.loginfo('Get Obj Cate: \n {}'.format(obj_cate))
    #         rospy.loginfo('Get Obj Pose: \n {}'.format(obj_pose.data))
    #         rospy.loginfo('Get Obj Scale: \n {}'.format(obj_scale))
    #         rospy.loginfo('Get Obj Scales: \n {}'.format(obj_scales))
    #     else:
    #         rospy.loginfo('WARNING! No Object detected! Please check again!')
    #     break

    obj_pose = (0.21940892934799194, -0.03242618590593338, -0.07573219388723373, 0.07833033800125122, 0.0431068055331707, -0.13841032981872559, 0.18415071070194244, -0.037015024572610855, -0.07020417600870132, -0.18632827699184418, -0.12361331284046173, 0.47620677947998047, 0.0, 0.0, 0.0, 1.0)
    obj_pose = np.array(obj_pose).reshape([4, 4])
    obj_pose = obj_pose[np.newaxis, :]

    obj_scales = (0.7724042534828186, 0.5262684226036072, 0.46758586168289185)
    obj_scales = np.array(obj_scales)
    obj_scales = obj_scales[np.newaxis, :]

    obj_scale = 0.234365314245224
    obj_scale = np.array(obj_scale).reshape([-1, 1])

    detector = Detector()
    boxes = []
    for i, (pose, s, scale) in enumerate(zip(obj_pose, obj_scale, obj_scales)):
        R = pose[:3, :3] * s
        T = pose[:3, 3].reshape(-1, 1)

        pred_sRT = np.concatenate((R, T), axis=1)
        pred_sRT = np.concatenate((pred_sRT, pose[3, :].reshape(-1, 4)), axis=0)

        # box = detector.get_3d_box(T, scale, R)
        box = draw_detections(pred_sRT[np.newaxis, :], scale[np.newaxis, :])
        box = box.reshape(3, 8).T

        box=box.transpose(1,0).ravel()
        boxes.append(box)

    while not rospy.is_shutdown():
        detector.display(boxes)
        rospy.loginfo('Show bbox')
        # break

    # for i, (pose, s, scale) in enumerate(zip(obj_pose, obj_scale, obj_scales)):
    #     R = pose[:3, :3]
    #     T = pose[:3, 3].reshape(-1, 1)

    #     RT = np.concatenate((R, T), axis=1)
    #     RT = np.concatenate((RT, pose[3, :].reshape(-1, 4)), axis=0)

    #     marker = Pose_Marker(RT)
    #     marker.pose_marker()